import os
from argparse import ArgumentParser

import torch

def get_new_name(org:str, newname:str, *, pre:str ='') -> str:
    name = org.split('_')
    if pre != '':
        name[0] = pre
    filename = name[0] + '_'  + newname + '_' + name[2]
    root, _ = os.path.splitext(filename)
    filename = root+".pth"
    return filename

if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument(
        '-s', '--src_dir', 
        default='/home/rslab/DLMatch/workspace/SOMatch/weights/default', 
        help='folder where the trained weights file saved'
    )
    parser.add_argument(
        '-w', '--weights_dir', 
        default='/home/rslab/DLMatch/workspace/SOMatch/weights/default', 
        help='folder to save the pth file'
    )
    parser.add_argument(
        '-f', '--pt_file', default="checkpoint_checkpoint_110.pt", 
        help='name of check point file'
    )
    parser.add_argument(
        '-p', '--prefix', default="default", help='perfix of pth file name'
    )
    args = parser.parse_args()

    src_dir = args.src_dir
    weights_dir = args.weights_dir
    weights_filename = args.pt_file
    prefix = args.prefix
    weights_path = os.path.join(src_dir, weights_filename)

    if not os.path.exists(weights_dir):
        os.mkdir(weights_dir)

    model_dict = torch.load(weights_path, map_location='cpu')

    for k in model_dict.keys():
        print(k)
        new_path = os.path.join(weights_dir, get_new_name(weights_filename, k, pre=prefix))
        torch.save(model_dict[k], new_path)

        # new_path = os.path.join(weights_dir, get_new_name(weights_filename, 'FtsA', pre=prefix))
        # torch.save(model_dict['FtsA'], new_path)
        # new_path = os.path.join(weights_dir, get_new_name(weights_filename, 'FtsB', pre=prefix))
        # torch.save(model_dict['FtsB'], new_path)
        # new_path = os.path.join(weights_dir, get_new_name(weights_filename, 'Fts', pre=prefix))
        # torch.save(model_dict['Fts'], new_path)

    